// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/expression/eval_ctx.h>
#include <pollux/expression/expr.h>
#include <pollux/expression/string_writer.h>
#include <pollux/expression/vector_function.h>
#include <pollux/expression/vector_writers.h>

namespace kumo::pollux::functions {
    namespace {
        /**
         * split(string, delimiter) -> array(varchar)
         *           Splits string on delimiter and returns an array.

         * split(string, delimiter, limit) -> array(varchar)
         *           Splits string on delimiter and returns an array of size at most
         limit.
         * The last element in the array will contain the remainder of the string,
         if such left.
         * limit must be a positive number.
         */

        class SplitFunction : public exec::VectorFunction {
        public:
            static std::vector<std::shared_ptr<exec::FunctionSignature> > signatures() {
                std::vector<std::shared_ptr<exec::FunctionSignature> > signatures;

                // varchar, varchar -> array(varchar)
                signatures.emplace_back(exec::FunctionSignatureBuilder()
                    .returnType("array(varchar)")
                    .argumentType("varchar")
                    .argumentType("varchar")
                    .build());

                // varchar, varchar, integer|bigint -> array(varchar)
                for (const auto &limitType: {"integer", "bigint"}) {
                    signatures.emplace_back(exec::FunctionSignatureBuilder()
                        .returnType("array(varchar)")
                        .argumentType("varchar")
                        .argumentType("varchar")
                        .argumentType(limitType)
                        .build());
                }
                return signatures;
            }

        private:
            /**
             * The function calls the 'typed' template function based on the type of the
             * 'limit' argument.
             */
            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr & /* outputType */,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                // If no 'limit' specified we just pick some int type.
                // Should we have two passes, with limit and without?
                TypeKind limitType =
                        (args.size() == 2) ? TypeKind::INTEGER : args[2]->type_kind();
                switch (limitType) {
                    case TypeKind::INTEGER:
                        applyTyped<int32_t>(rows, args, context, result);
                        return;
                    case TypeKind::BIGINT:
                        applyTyped<int64_t>(rows, args, context, result);
                        return;
                    default:
                        POLLUX_FAIL(
                            "Unsupported type for 'limit' argument of 'split' function: {}",
                            mapTypeKindToName(limitType));
                }
            }

            template<typename I>
            void applyTyped(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                exec::EvalCtx &context,
                VectorPtr &result) const {
                // Get the decoded vectors out of arguments.
                const bool noLimit = (args.size() == 2);
                exec::DecodedArgs decodedArgs(rows, args, context);
                DecodedVector *strings = decodedArgs.at(0);
                DecodedVector *delims = decodedArgs.at(1);
                DecodedVector *limits = noLimit ? nullptr : decodedArgs.at(2);

                BaseVector::ensure_writable(rows, ARRAY(VARCHAR()), context.pool(), result);
                exec::VectorWriter<Array<Varchar> > resultWriter;
                resultWriter.init(*result->as<ArrayVector>());

                // Optimization for the (flat, const, const) case.
                if (strings->isIdentityMapping() and delims->isConstantMapping() and
                    (noLimit or limits->isConstantMapping())) {
                    const auto *rawStrings = strings->data<StringView>();
                    const auto delim = delims->value_at<StringView>(0);

                    if (noLimit) {
                        const I limit = std::numeric_limits<I>::max();
                        rows.applyToSelected([&](vector_size_t row) {
                            applyInner<false, I>(
                                rawStrings[row], delim, limit, row, resultWriter);
                        });
                    } else {
                        const I limit = limits->value_at<I>(0);
                        // Limit must be positive.
                        if (limit > 0) {
                            rows.applyToSelected([&](vector_size_t row) {
                                applyInner<true, I>(
                                    rawStrings[row], delim, limit, row, resultWriter);
                            });
                        } else {
                            auto pex = std::make_exception_ptr(
                                std::invalid_argument("Limit must be positive"));
                            context.setErrors(rows, pex);
                        }
                    }
                } else {
                    // The rest of the cases are handled through this general path and no
                    // direct access.
                    applyDecoded<I>(rows, context, strings, delims, limits, resultWriter);
                }

                resultWriter.finish();

                // Ensure that our result elements vector uses the same string buffer as
                // the input vector of strings.
                result->as<ArrayVector>()
                        ->elements()
                        ->as<FlatVector<StringView> >()
                        ->acquireSharedStringBuffers(strings->base());
            }

            template<typename I>
            void applyDecoded(
                const SelectivityVector &rows,
                exec::EvalCtx &context,
                DecodedVector *strings,
                DecodedVector *delims,
                DecodedVector *limits,
                exec::VectorWriter<Array<Varchar> > &resultWriter) const {
                if (limits == nullptr) {
                    auto limit = std::numeric_limits<I>::max();
                    rows.applyToSelected([&](vector_size_t row) {
                        applyInner<false, I>(
                            strings->value_at<StringView>(row),
                            delims->value_at<StringView>(row),
                            limit,
                            row,
                            resultWriter);
                    });
                } else {
                    rows.applyToSelected([&](vector_size_t row) {
                        const I limit = limits->value_at<I>(row);
                        if (limit > 0) {
                            applyInner<true, I>(
                                strings->value_at<StringView>(row),
                                delims->value_at<StringView>(row),
                                limit,
                                row,
                                resultWriter);
                        } else {
                            auto pex = std::make_exception_ptr(
                                std::invalid_argument("Limit must be positive"));
                            context.setError(row, pex);
                        }
                    });
                }
            }

            /**
             * The inner most kernel of the vector operations for 'split'.
             */
            template<bool hasLimit = false, typename I>
            inline void applyInner(
                StringView input,
                const StringView delim,
                I limit,
                vector_size_t row,
                exec::VectorWriter<Array<Varchar> > &resultWriter) const {
                // Add new array (for the new row) to our array vector.
                resultWriter.setOffset(row);
                auto &arrayWriter = resultWriter.current();

                // Trivial case of converting string to array with 1 element.
                if (hasLimit and limit == 1) {
                    arrayWriter.add_item().setNoCopy(input);
                    resultWriter.commit();
                    return;
                }

                // We walk through our input cutting off the pieces using the delimiter and
                // adding them to the elements vector, until we reached the end of the
                // string or the limit.
                int32_t addedElements{0};
                std::string_view sinput(input.data(), input.size());
                const std::string_view sdelim(delim.data(), delim.size());
                while (true) {
                    // Find the byte of the 1st delimiter.
                    auto byteIndex = sinput.find(sdelim, 0);

                    // Special case for empty delimiters. Split character by character with an
                    // empty string at the end.
                    if (sdelim.empty()) {
                        byteIndex++;
                    }

                    // Delimiter is not found, leave the loop.
                    if (byteIndex > sinput.size()) {
                        break;
                    }

                    // Add the new element, we've split
                    arrayWriter.add_item().setNoCopy(StringView(sinput.data(), byteIndex));

                    // Advance input by the size of the element + delimiter.
                    // Note: should we add 'advance' method?
                    const auto advanceBytes = byteIndex + sdelim.size();
                    sinput = std::string_view(
                        sinput.data() + advanceBytes, sinput.size() - advanceBytes);

                    if (hasLimit) {
                        ++addedElements;
                        // If the next element should be the last, leave the loop.
                        if (addedElements + 1 == limit) {
                            break;
                        }
                    }
                }

                // Add the rest of the string and we are done.
                // Note, that the rest of the string can be empty - we still add it.
                arrayWriter.add_item().setNoCopy(StringView(sinput.data(), sinput.size()));
                resultWriter.commit();
            }
        };
    } // namespace

    POLLUX_DECLARE_VECTOR_FUNCTION(
        udf_split,
        SplitFunction::signatures(),
        std::make_unique<SplitFunction>());
} // namespace kumo::pollux::functions
